#!/usr/bin/env python3

import traceback
import trimesh
import argparse
import os
import sys

import numpy as np
import re
def convert_glb_to_obj(glb_path, output_path=None, suppress_print=False):
    if output_path is None:
        output_path = glb_path.replace('.glb', '.obj')
    
    if not suppress_print:
        print(f"Loading GLB: {glb_path}")
    scene = trimesh.load(glb_path)
    
    if hasattr(scene, 'geometry'):
        mesh = list(scene.geometry.values())[0]
    else:
        mesh = scene
    
    if not suppress_print:
        print(f"Mesh vertices: {len(mesh.vertices)}")
        print(f"Mesh faces: {len(mesh.faces)}")
        print(f"Has vertex colors: {hasattr(mesh.visual, 'vertex_colors') and mesh.visual.vertex_colors is not None}")
        print(f"Has material: {hasattr(mesh.visual, 'material')}")
        
        if hasattr(mesh.visual, 'material'):
            print(f"Material type: {type(mesh.visual.material)}")
            if hasattr(mesh.visual.material, 'main_color'):
                print(f"Main color: {mesh.visual.material.main_color}")
            if hasattr(mesh.visual.material, 'baseColorFactor'):
                print(f"Base color factor: {mesh.visual.material.baseColorFactor}")
    
    if not suppress_print:
        print(f"Exporting to OBJ: {output_path}")
    mesh.export(output_path)
    
    glb_basename = os.path.basename(glb_path)
    match = re.search(r'base(\d+)\.glb', glb_basename)
    if match:
        number = match.group(1)
        obj_dir = os.path.dirname(output_path)
        old_mtl_path = os.path.join(obj_dir, 'material.mtl')
        new_mtl_path = os.path.join(obj_dir, f'base{number}.mtl')
        
        if os.path.exists(old_mtl_path):
            os.rename(old_mtl_path, new_mtl_path)
            if not suppress_print:
                print(f"Renamed material.mtl to base{number}.mtl")
            
            with open(output_path, 'r') as f:
                obj_content = f.read()
            
            obj_content = obj_content.replace('material.mtl', f'base{number}.mtl')
            
            with open(output_path, 'w') as f:
                f.write(obj_content)
            
            if not suppress_print:
                print(f"Updated OBJ file to reference base{number}.mtl")
    
    if not suppress_print:
        print(f"Conversion completed: {output_path}")
    return output_path

def test_obj_with_current_code(obj_path):
    print(f"\nTesting OBJ file with current code: {obj_path}")
    mesh = trimesh.load(obj_path)
    
    print("Testing basic mesh operations...")
    
    print(f"1. mesh.vertices shape: {mesh.vertices.shape}")
    print(f"2. mesh.vertex_normals shape: {mesh.vertex_normals.shape}")
    print(f"3. mesh.faces shape: {mesh.faces.shape}")
    
    print("4. Testing mesh.copy()...")
    mesh_copy = mesh.copy()
    assert mesh_copy.vertices.shape == mesh.vertices.shape
    print("   ✓ mesh.copy() works")
    
    print("5. Testing mesh.vertices operations...")
    max_xyz = mesh.vertices.max(axis=0)
    min_xyz = mesh.vertices.min(axis=0)
    model_center = (min_xyz + max_xyz) / 2
    print(f"   max_xyz: {max_xyz}")
    print(f"   min_xyz: {min_xyz}")
    print(f"   model_center: {model_center}")
    
    print("6. Testing mesh vertex modification...")
    mesh_centered = mesh.copy()
    mesh_centered.vertices = mesh_centered.vertices - model_center.reshape(1, 3)
    print(f"   Centered vertices shape: {mesh_centered.vertices.shape}")
    
    print("7. Testing trimesh.bounds.oriented_bounds...")
    to_origin, extents = trimesh.bounds.oriented_bounds(mesh)
    bbox = np.stack([-extents/2, extents/2], axis=0).reshape(2, 3)
    print(f"   to_origin shape: {to_origin.shape}")
    print(f"   extents: {extents}")
    print(f"   bbox shape: {bbox.shape}")
    
    print("8. Testing mesh.export...")
    test_export_path = obj_path.replace('.obj', '_test_export.obj')
    mesh.export(test_export_path)
    assert os.path.exists(test_export_path)
    print(f"   ✓ Export successful: {test_export_path}")
    
    print("9. Testing mesh.apply_transform...")
    test_pose = np.eye(4)
    test_pose[:3, 3] = [0.1, 0.2, 0.3]
    mesh_transformed = mesh.copy()
    mesh_transformed.apply_transform(test_pose)
    print(f"   ✓ Transform applied, new center: {mesh_transformed.vertices.mean(axis=0)}")
    
    print("10. Testing FoundationPose initialization...")
    try:
        scorer = ScorePredictor()
        refiner = PoseRefinePredictor()
        glctx = dr.RasterizeCudaContext()
        
        est = FoundationPose(
            model_pts=mesh.vertices, 
            model_normals=mesh.vertex_normals, 
            mesh=mesh, 
            scorer=scorer, 
            refiner=refiner, 
            debug_dir='/tmp/test_debug', 
            debug=0, 
            glctx=glctx
        )
        print("   ✓ FoundationPose initialization successful")
        
        print("11. Testing est.reset_object...")
        est.reset_object(
            model_pts=mesh.vertices, 
            model_normals=mesh.vertex_normals, 
            mesh=mesh
        )
        print("   ✓ reset_object successful")
        
    except Exception as e:
        print(f"   ✗ FoundationPose test failed: {traceback.format_exc()}")
    
    print("12. Testing mesh visual properties...")
    if hasattr(mesh.visual, 'vertex_colors') and mesh.visual.vertex_colors is not None:
        print(f"   Vertex colors shape: {mesh.visual.vertex_colors.shape}")
    if hasattr(mesh.visual, 'material'):
        print(f"   Material type: {type(mesh.visual.material)}")
    
    print("13. Testing mesh diameter computation...")
    try:
        from Utils import compute_mesh_diameter
        diameter = compute_mesh_diameter(model_pts=mesh.vertices, n_sample=1000)
        print(f"   Mesh diameter: {diameter}")
    except Exception as e:
        print(f"   ✗ Diameter computation failed: {e}")
    
    print("14. Testing mesh tensors creation...")
    try:
        from Utils import make_mesh_tensors
        mesh_tensors = make_mesh_tensors(mesh)
        print(f"   ✓ Mesh tensors created with keys: {list(mesh_tensors.keys())}")
    except Exception as e:
        print(f"   ✗ Mesh tensors creation failed: {e}")
    
    print("\nAll mesh operation tests completed!")
    
    if os.path.exists(test_export_path):
        os.remove(test_export_path)
        print(f"Cleaned up test file: {test_export_path}")



if __name__ == '__main__':
    # from estimater import *
    # from datareader import *
    parser = argparse.ArgumentParser()
    parser.add_argument('glb_file', type=str, help='Path to GLB file')
    parser.add_argument('--output', type=str, default=None, help='Output OBJ path (optional)')
    args = parser.parse_args()
    
    if not os.path.exists(args.glb_file):
        print(f"Error: GLB file not found: {args.glb_file}")
        sys.exit(1)
    
    try:
        obj_path = convert_glb_to_obj(args.glb_file, args.output)
        
        # test_obj_with_current_code(obj_path)
            
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
